import os
import cv2


def remove_non_image_files(dataset_path):
    """
    删除数据集中非图像文件
    :param dataset_path: 数据集路径
    """
    valid_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
    for filename in os.listdir(dataset_path):
        if not filename.lower().endswith(valid_extensions):
            file_path = os.path.join(dataset_path, filename)
            os.remove(file_path)
            print(f"删除非图像文件: {filename}")


def convert_images_to_jpg(dataset_path, output_path):
    """
    将数据集中所有图像读取并存储为jpg格式
    :param dataset_path: 数据集路径
    :param output_path: 输出路径
    """
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    for filename in os.listdir(dataset_path):
        if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            file_path = os.path.join(dataset_path, filename)
            image = cv2.imread(file_path)  # 用opencv读取图像

            if image is not None:
                output_file = os.path.join(output_path, f"{os.path.splitext(filename)[0]}.jpg")  # 确定图像保存路径
                cv2.imwrite(output_file, image)
                print(f"转换并保存图像: {output_file}")


if __name__ == "__main__":
    dataset_path = 'path/to/your/dataset'  # 替换为数据集路径
    output_path = 'path/to/save/jpg_images'  # 替换为输出路径

    # 举例:
    # dataset_path = './VOC2028_short/images1'  # 替换为数据集路径
    # output_path = './VOC2028_short/images'  # 替换为输出路径

    remove_non_image_files(dataset_path)
    convert_images_to_jpg(dataset_path, output_path)